import numpy as np
#import matplotlib.pyplot as plt
from sklearn.decomposition import SparseCoder
import timeit

COMPONENT_COUNT = 350
PATCH_COUNT = 400000
EPOCH_COUNT = 16
SPARSE_CODE_COUNT = 2800
SPARSITY_PARAMETER = 4.0
NON_NEGATIVE = True

start = timeit.default_timer()

sortedFile = 'codes/codesSortednn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
codesFile = 'codes/codesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'

#sortedFile = 'texcodesSortednn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) +  'patches' + str(PATCH_COUNT) + '.npy'
#codesFile = 'texcodesnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + '.npy'

pcaTransformed = np.load('pcaTransformed.npy')[:, :COMPONENT_COUNT]

inputFile = 'basis/basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(SPARSE_CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
print(inputFile)
print(codesFile)

basis = np.load(inputFile)

sc = SparseCoder(basis, transform_algorithm = 'lasso_cd', transform_alpha = SPARSITY_PARAMETER, positive_code = NON_NEGATIVE)

codes = sc.transform(pcaTransformed)

codesSorted = np.argsort(codes, axis = 0)

np.save(sortedFile, codesSorted)
np.save(codesFile, codes)

stop = timeit.default_timer()

print('Time: ', stop - start)
